// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <nebula/core/finger_printable.h>
#include <nebula/types/type_fwd.h>
#include <turbo/meta/compare.h>
#include <nebula/util/key_value_metadata.h>

namespace nebula {

    /// \brief The combination of a field name and data type, with optional metadata
    ///
    /// Fields are used to describe the individual constituents of a
    /// nested DataType or a Schema.
    ///
    /// A field's metadata is represented by a KeyValueMetadata instance,
    /// which holds arbitrary key-value pairs.
    class TURBO_EXPORT Field : public detail::FingerPrintAble,
                               public turbo::EqualityComparable<Field> {
    public:
        Field(std::string name, std::shared_ptr<DataType> type, bool nullable = true,
              std::shared_ptr<const KeyValueMetadata> metadata = nullptr)
                : detail::FingerPrintAble(),
                  name_(std::move(name)),
                  type_(std::move(type)),
                  nullable_(nullable),
                  metadata_(std::move(metadata)) {}

        ~Field() override = default;

        /// \brief Return the field's attached metadata
        std::shared_ptr<const KeyValueMetadata> metadata() const { return metadata_; }

        /// \brief Return whether the field has non-empty metadata
        bool has_metadata() const;

        /// \brief Return a copy of this field with the given metadata attached to it
        std::shared_ptr<Field> with_metadata(
                const std::shared_ptr<const KeyValueMetadata> &metadata) const;

        /// \brief EXPERIMENTAL: Return a copy of this field with the given metadata
        /// merged with existing metadata (any colliding keys will be overridden by
        /// the passed metadata)
        std::shared_ptr<Field> with_merged_metadata(
                const std::shared_ptr<const KeyValueMetadata> &metadata) const;

        /// \brief Return a copy of this field without any metadata attached to it
        std::shared_ptr<Field> remove_metadata() const;

        /// \brief Return a copy of this field with the replaced type.
        std::shared_ptr<Field> with_type(const std::shared_ptr<DataType> &type) const;

        /// \brief Return a copy of this field with the replaced name.
        std::shared_ptr<Field> with_name(const std::string &name) const;

        /// \brief Return a copy of this field with the replaced nullability.
        std::shared_ptr<Field> with_nullable(bool nullable) const;

        /// \brief Options that control the behavior of `merge_with`.
        /// Options are to be added to allow type conversions, including integer
        /// widening, promotion from integer to float, or conversion to or from boolean.
        struct TURBO_EXPORT MergeOptions {
            /// If true, a Field of NullType can be unified with a Field of another type.
            /// The unified field will be of the other type and become nullable.
            /// Nullability will be promoted to the looser option (nullable if one is not
            /// nullable).
            bool promote_nullability = true;

            /// Allow a decimal to be unified with another decimal of the same
            /// width, adjusting scale and precision as appropriate. May fail
            /// if the adjustment is not possible.
            bool promote_decimal = false;

            /// Allow a decimal to be promoted to a float. The float type will
            /// not itself be promoted (e.g. Decimal128 + Float32 = Float32).
            bool promote_decimal_to_float = false;

            /// Allow an integer to be promoted to a decimal.
            ///
            /// May fail if the decimal has insufficient precision to
            /// accommodate the integer (see promote_numeric_width).
            bool promote_integer_to_decimal = false;

            /// Allow an integer of a given bit width to be promoted to a
            /// float; the result will be a float of an equal or greater bit
            /// width to both of the inputs. Examples:
            ///  - int8 + float32 = float32
            ///  - int32 + float32 = float64
            ///  - int32 + float64 = float64
            /// Because an int32 cannot always be represented exactly in the
            /// 24 bits of a float32 mantissa.
            bool promote_integer_to_float = false;

            /// Allow an unsigned integer of a given bit width to be promoted
            /// to a signed integer that fits into the signed type:
            /// uint + int16 = int16
            /// When widening is needed, set promote_numeric_width to true:
            /// uint16 + int16 = int32
            bool promote_integer_sign = false;

            /// Allow an integer, float, or decimal of a given bit width to be
            /// promoted to an equivalent type of a greater bit width.
            bool promote_numeric_width = false;

            /// Allow strings to be promoted to binary types. Promotion of fixed size
            /// binary types to variable sized formats, and binary to large binary,
            /// and string to large string.
            bool promote_binary = false;

            /// Second to millisecond, Time32 to Time64, Time32(SECOND) to Time32(MILLI), etc
            bool promote_temporal_unit = false;

            /// Allow promotion from a list to a large-list and from a fixed-size list to a
            /// variable sized list
            bool promote_list = false;

            /// unify dictionary index types and dictionary value types.
            bool promote_dictionary = false;

            /// Allow merging ordered and non-ordered dictionaries.
            /// The result will be ordered if and only if both inputs
            /// are ordered.
            bool promote_dictionary_ordered = false;

            /// get default options. Only NullType will be merged with other types.
            static MergeOptions defaults() { return {}; }

            /// get permissive options. All options are enabled, except
            /// promote_dictionary_ordered.
            static MergeOptions permissive();

            /// get a human-readable representation of the options.
            [[nodiscard]] std::string to_string() const;
        };

        /// \brief Merge the current field with a field of the same name.
        ///
        /// The two fields must be compatible, i.e:
        ///   - have the same name
        ///   - have the same type, or of compatible types according to `options`.
        ///
        /// The metadata of the current field is preserved; the metadata of the other
        /// field is discarded.
        turbo::Result<std::shared_ptr<Field>> merge_with(
                const Field &other, MergeOptions options = MergeOptions::defaults()) const;

        turbo::Result<std::shared_ptr<Field>> merge_with(
                const std::shared_ptr<Field> &other,
                MergeOptions options = MergeOptions::defaults()) const;

        FieldVector Flatten() const;

        /// \brief Indicate if fields are equals.
        ///
        /// \param[in] other field to check equality with.
        /// \param[in] check_metadata controls if it should check for metadata
        ///            equality.
        ///
        /// \return true if fields are equal, false otherwise.
        bool equals(const Field &other, bool check_metadata = false) const;

        bool equals(const std::shared_ptr<Field> &other, bool check_metadata = false) const;

        /// \brief Indicate if fields are compatibles.
        ///
        /// See the criteria of merge_with.
        ///
        /// \return true if fields are compatible, false otherwise.
        bool is_compatible_with(const Field &other) const;

        bool is_compatible_with(const std::shared_ptr<Field> &other) const;

        /// \brief Return a string representation ot the field
        /// \param[in] show_metadata when true, if KeyValueMetadata is non-empty,
        /// print keys and values in the output
        std::string to_string(bool show_metadata = false) const;

        /// \brief Return the field name
        const std::string &name() const { return name_; }

        /// \brief Return the field data type
        const std::shared_ptr<DataType> &type() const { return type_; }

        /// \brief Return whether the field is nullable
        bool nullable() const { return nullable_; }

        std::shared_ptr<Field> deep_copy() const;

    private:
        std::string compute_fingerprint() const override;

        std::string compute_metadata_fingerprint() const override;

        // Field name
        std::string name_;

        // The field's data type
        std::shared_ptr<DataType> type_;

        // Fields can be nullable
        bool nullable_;

        // The field's metadata, if any
        std::shared_ptr<const KeyValueMetadata> metadata_;

        TURBO_DISALLOW_COPY_AND_ASSIGN(Field);
    };

    TURBO_EXPORT void print_to(const Field &field, std::ostream *os);

    // ----------------------------------------------------------------------
    // FieldRef

    /// \class FieldPath
    ///
    /// Represents a path to a nested field using indices of child fields.
    /// For example, given indices {5, 9, 3} the field would be retrieved with
    /// schema->field(5)->type()->field(9)->type()->field(3)
    ///
    /// Attempting to retrieve a child field using a FieldPath which is not valid for
    /// a given schema will raise an error. Invalid FieldPaths include:
    /// - an index is out of range
    /// - the path is empty (note: a default constructed FieldPath will be empty)
    ///
    /// FieldPaths provide a number of accessors for drilling down to potentially nested
    /// children. They are overloaded for convenience to support Schema (returns a field),
    /// DataType (returns a child field), Field (returns a child field of this field's type)
    /// Array (returns a child array), RecordBatch (returns a column).
    class TURBO_EXPORT FieldPath {
    public:
        FieldPath() = default;

        FieldPath(std::vector<int> indices)  // NOLINT runtime/explicit
                : indices_(std::move(indices)) {}

        FieldPath(std::initializer_list<int> indices)  // NOLINT runtime/explicit
                : indices_(indices) {}

        [[nodiscard]] std::string to_string() const;

        [[nodiscard]] size_t hash() const;

        struct Hash {
            size_t operator()(const FieldPath &path) const { return path.hash(); }
        };

        [[nodiscard]] bool empty() const { return indices_.empty(); }

        bool operator==(const FieldPath &other) const { return indices() == other.indices(); }

        bool operator!=(const FieldPath &other) const { return indices() != other.indices(); }

        [[nodiscard]] const std::vector<int> &indices() const { return indices_; }

        int operator[](size_t i) const { return indices_[i]; }

        [[nodiscard]] std::vector<int>::const_iterator begin() const { return indices_.begin(); }

        [[nodiscard]] std::vector<int>::const_iterator end() const { return indices_.end(); }

        /// \brief Retrieve the referenced child Field from a Schema, Field, or DataType
        [[nodiscard]] turbo::Result<std::shared_ptr<Field>> get(const Schema &schema) const;

        [[nodiscard]] turbo::Result<std::shared_ptr<Field>> get(const Field &field) const;

        [[nodiscard]] turbo::Result<std::shared_ptr<Field>> get(const DataType &type) const;

        [[nodiscard]] turbo::Result<std::shared_ptr<Field>> get(const FieldVector &fields) const;

        static turbo::Result<std::shared_ptr<Schema>> get_all(const Schema &schema,
                                                             const std::vector<FieldPath> &paths);

        /// \brief Retrieve the referenced column from a RecordBatch or Table
        [[nodiscard]] turbo::Result<std::shared_ptr<Array>> get(const RecordBatch &batch) const;

        [[nodiscard]] turbo::Result<std::shared_ptr<ChunkedArray>> get(const Table &table) const;

        /// \brief Retrieve the referenced child from an Array or ArrayData
        [[nodiscard]] turbo::Result<std::shared_ptr<Array>> get(const Array &array) const;

        [[nodiscard]] turbo::Result<std::shared_ptr<ArrayData>> get(const ArrayData &data) const;

        /// \brief Retrieve the referenced child from a ChunkedArray
        [[nodiscard]] turbo::Result<std::shared_ptr<ChunkedArray>> get(const ChunkedArray &chunked_array) const;

        /// \brief Retrieve the referenced child/column from an Array, ArrayData, ChunkedArray,
        /// RecordBatch, or Table
        ///
        /// Unlike `FieldPath::get`, these variants are not zero-copy and the retrieved child's
        /// null bitmap is ANDed with its ancestors'
        turbo::Result<std::shared_ptr<Array>> get_flattened(const Array &array,
                                                           MemoryPool *pool = nullptr) const;

        turbo::Result<std::shared_ptr<ArrayData>> get_flattened(const ArrayData &data,
                                                               MemoryPool *pool = nullptr) const;

        turbo::Result<std::shared_ptr<ChunkedArray>> get_flattened(const ChunkedArray &chunked_array,
                                                                  MemoryPool *pool = nullptr) const;

        turbo::Result<std::shared_ptr<Array>> get_flattened(const RecordBatch &batch,
                                                           MemoryPool *pool = nullptr) const;

        turbo::Result<std::shared_ptr<ChunkedArray>> get_flattened(const Table &table,
                                                                  MemoryPool *pool = nullptr) const;

    private:
        std::vector<int> indices_;
    };

    /// \class FieldRef
    /// \brief Descriptor of a (potentially nested) field within a schema.
    ///
    /// Unlike FieldPath (which exclusively uses indices of child fields), FieldRef may
    /// reference a field by name. It is intended to replace parameters like `int field_index`
    /// and `const std::string& field_name`; it can be implicitly constructed from either a
    /// field index or a name.
    ///
    /// Nested fields can be referenced as well. Given
    ///     schema({field("a", STRUCT({field("n", null())})), field("b", int32())})
    ///
    /// the following all indicate the nested field named "n":
    ///     FieldRef ref1(0, 0);
    ///     FieldRef ref2("a", 0);
    ///     FieldRef ref3("a", "n");
    ///     FieldRef ref4(0, "n");
    ///     TURBO_MOVE_OR_RAISE(FieldRef ref5,
    ///                           FieldRef::from_dot_path(".a[0]"));
    ///
    /// FieldPaths matching a FieldRef are retrieved using the member function find_all.
    /// Multiple matches are possible because field names may be duplicated within a schema.
    /// For example:
    ///     Schema a_is_ambiguous({field("a", int32()), field("a", float32())});
    ///     auto matches = FieldRef("a").find_all(a_is_ambiguous);
    ///     assert(matches.size() == 2);
    ///     assert(matches[0].get(a_is_ambiguous)->equals(a_is_ambiguous.field(0)));
    ///     assert(matches[1].get(a_is_ambiguous)->equals(a_is_ambiguous.field(1)));
    ///
    /// Convenience accessors are available which raise a helpful error if the field is not
    /// found or ambiguous, and for immediately calling FieldPath::get to retrieve any
    /// matching children:
    ///     auto maybe_match = FieldRef("struct", "field_i32").find_one_or_none(schema);
    ///     auto maybe_column = FieldRef("struct", "field_i32").get_one(some_table);
    class TURBO_EXPORT FieldRef : public turbo::EqualityComparable<FieldRef> {
    public:
        FieldRef() = default;

        /// Construct a FieldRef using a string of indices. The reference will be retrieved as:
        /// schema.fields[self.indices[0]].type.fields[self.indices[1]] ...
        ///
        /// Empty indices are not valid.
        FieldRef(FieldPath indices);  // NOLINT runtime/explicit

        /// Construct a by-name FieldRef. Multiple fields may match a by-name FieldRef:
        /// [f for f in schema.fields where f.name == self.name]
        FieldRef(std::string name) : impl_(std::move(name)) {}    // NOLINT runtime/explicit
        FieldRef(const char *name) : impl_(std::string(name)) {}  // NOLINT runtime/explicit

        /// Equivalent to a single index string of indices.
        FieldRef(int index) : impl_(FieldPath({index})) {}  // NOLINT runtime/explicit

        /// Construct a nested FieldRef.
        explicit FieldRef(std::vector<FieldRef> refs) { Flatten(std::move(refs)); }

        /// Convenience constructor for nested FieldRefs: each argument will be used to
        /// construct a FieldRef
        template<typename A0, typename A1, typename... A>
        FieldRef(A0 &&a0, A1 &&a1, A &&... a) {
            Flatten({// cpplint thinks the following are constructor decls
                            FieldRef(std::forward<A0>(a0)),     // NOLINT runtime/explicit
                            FieldRef(std::forward<A1>(a1)),     // NOLINT runtime/explicit
                            FieldRef(std::forward<A>(a))...});  // NOLINT runtime/explicit
        }

        /// Parse a dot path into a FieldRef.
        ///
        /// dot_path = '.' name
        ///          | '[' digit+ ']'
        ///          | dot_path+
        ///
        /// Examples:
        ///   ".alpha" => FieldRef("alpha")
        ///   "[2]" => FieldRef(2)
        ///   ".beta[3]" => FieldRef("beta", 3)
        ///   "[5].gamma.delta[7]" => FieldRef(5, "gamma", "delta", 7)
        ///   ".hello world" => FieldRef("hello world")
        ///   R"(.\[y\]\\tho\.\)" => FieldRef(R"([y]\tho.\)")
        ///
        /// Note: When parsing a name, a '\' preceding any other character will be dropped from
        /// the resulting name. Therefore if a name must contain the characters '.', '\', or '['
        /// those must be escaped with a preceding '\'.
        static turbo::Result<FieldRef> from_dot_path(const std::string &dot_path);

        [[nodiscard]] std::string to_dot_path() const;

        [[nodiscard]] bool equals(const FieldRef &other) const { return impl_ == other.impl_; }

        [[nodiscard]] std::string to_string() const;

        [[nodiscard]] size_t hash() const;

        struct Hash {
            size_t operator()(const FieldRef &ref) const { return ref.hash(); }
        };

        explicit operator bool() const { return equals(FieldPath{}); }

        bool operator!() const { return !equals(FieldPath{}); }

        [[nodiscard]] bool is_field_path() const { return std::holds_alternative<FieldPath>(impl_); }

        [[nodiscard]] bool is_name() const { return std::holds_alternative<std::string>(impl_); }

        [[nodiscard]] bool is_nested() const {
            if (is_name()) return false;
            if (is_field_path()) return std::get<FieldPath>(impl_).indices().size() > 1;
            return true;
        }

        /// \brief Return true if this ref is a name or a nested sequence of only names
        ///
        /// Useful for determining if iteration is possible without recursion or inner loops
        [[nodiscard]]  bool is_name_sequence() const {
            if (is_name()) return true;
            if (const auto *nested = nested_refs()) {
                for (const auto &ref: *nested) {
                    if (!ref.is_name()) return false;
                }
                return !nested->empty();
            }
            return false;
        }

        [[nodiscard]] const FieldPath *field_path() const {
            return is_field_path() ? &std::get<FieldPath>(impl_) : nullptr;
        }

        [[nodiscard]] const std::string *name() const {
            return is_name() ? &std::get<std::string>(impl_) : nullptr;
        }

        [[nodiscard]] const std::vector<FieldRef> *nested_refs() const {
            return std::holds_alternative<std::vector<FieldRef>>(impl_)
                   ? &std::get<std::vector<FieldRef>>(impl_)
                   : nullptr;
        }

        /// \brief Retrieve FieldPath of every child field which matches this FieldRef.
        [[nodiscard]] std::vector<FieldPath> find_all(const Schema &schema) const;

        [[nodiscard]] std::vector<FieldPath> find_all(const Field &field) const;

        [[nodiscard]] std::vector<FieldPath> find_all(const DataType &type) const;

        [[nodiscard]] std::vector<FieldPath> find_all(const FieldVector &fields) const;

        /// \brief Convenience function which applies find_all to arg's type or schema.
        [[nodiscard]] std::vector<FieldPath> find_all(const ArrayData &array) const;

        [[nodiscard]] std::vector<FieldPath> find_all(const Array &array) const;

        [[nodiscard]] std::vector<FieldPath> find_all(const ChunkedArray &chunked_array) const;

        [[nodiscard]] std::vector<FieldPath> find_all(const RecordBatch &batch) const;

        [[nodiscard]] std::vector<FieldPath> find_all(const Table &table) const;

        /// \brief Convenience function: raise an error if matches is empty.
        template<typename T>
        turbo::Status check_non_empty(const std::vector<FieldPath> &matches, const T &root) const {
            if (matches.empty()) {
                return turbo::invalid_argument_error("No match for ", to_string(), " in ", root.to_string());
            }
            return turbo::OkStatus();
        }

        /// \brief Convenience function: raise an error if matches contains multiple FieldPaths.
        template<typename T>
        turbo::Status check_non_multiple(const std::vector<FieldPath> &matches, const T &root) const {
            if (matches.size() > 1) {
                return turbo::invalid_argument_error("Multiple matches for ", to_string(), " in ",
                                                     root.to_string());
            }
            return turbo::OkStatus();
        }

        /// \brief Retrieve FieldPath of a single child field which matches this
        /// FieldRef. Emit an error if none or multiple match.
        template<typename T>
        turbo::Result<FieldPath> find_one(const T &root) const {
            auto matches = find_all(root);
            TURBO_RETURN_NOT_OK(check_non_empty(matches, root));
            TURBO_RETURN_NOT_OK(check_non_multiple(matches, root));
            return std::move(matches[0]);
        }

        /// \brief Retrieve FieldPath of a single child field which matches this
        /// FieldRef. Emit an error if multiple match. An empty (invalid) FieldPath
        /// will be returned if none match.
        template<typename T>
        turbo::Result<FieldPath> find_one_or_none(const T &root) const {
            auto matches = find_all(root);
            TURBO_RETURN_NOT_OK(check_non_multiple(matches, root));
            if (matches.empty()) {
                return FieldPath();
            }
            return std::move(matches[0]);
        }

        template<typename T>
        using get_type = decltype(std::declval<FieldPath>().get(std::declval<T>()).value_or_die());

        /// \brief Get all children matching this FieldRef.
        template<typename T>
        std::vector<get_type<T>> get_all(const T &root) const {
            std::vector<get_type<T>> out;
            for (const auto &match: find_all(root)) {
                out.push_back(match.get(root).value_or_die());
            }
            return out;
        }

        /// \brief Get all children matching this FieldRef.
        ///
        /// Unlike `FieldRef::get_all`, this variant is not zero-copy and the retrieved
        /// children's null bitmaps are ANDed with their ancestors'
        template<typename T>
        turbo::Result<std::vector<get_type<T>>> get_all_flattened(const T &root,
                                                               MemoryPool *pool = nullptr) const {
            std::vector<get_type<T>> out;
            for (const auto &match: find_all(root)) {
                TURBO_MOVE_OR_RAISE(auto child, match.get_flattened(root, pool));
                out.push_back(std::move(child));
            }
            return out;
        }

        /// \brief Get the single child matching this FieldRef.
        /// Emit an error if none or multiple match.
        template<typename T>
        turbo::Result<get_type<T>> get_one(const T &root) const {
            TURBO_MOVE_OR_RAISE(auto match, find_one(root));
            return match.get(root).value_or_die();
        }

        /// \brief Get the single child matching this FieldRef.
        ///
        /// Unlike `FieldRef::get_one`, this variant is not zero-copy and the retrieved
        /// child's null bitmap is ANDed with its ancestors'
        template<typename T>
        turbo::Result<get_type<T>> get_one_flattened(const T &root, MemoryPool *pool = nullptr) const {
            TURBO_MOVE_OR_RAISE(auto match, find_one(root));
            return match.get_flattened(root, pool);
        }

        /// \brief Get the single child matching this FieldRef.
        /// Return nullptr if none match, emit an error if multiple match.
        template<typename T>
        turbo::Result<get_type<T>> get_one_or_none(const T &root) const {
            TURBO_MOVE_OR_RAISE(auto match, find_one_or_none(root));
            if (match.empty()) {
                return static_cast<get_type<T>>(nullptr);
            }
            return match.get(root).value_or_die();
        }

        /// \brief Get the single child matching this FieldRef.
        ///
        /// Return nullptr if none match, emit an error if multiple match.
        /// Unlike `FieldRef::get_one_or_none`, this variant is not zero-copy and the
        /// retrieved child's null bitmap is ANDed with its ancestors'
        template<typename T>
        turbo::Result<get_type<T>> get_one_or_none_flattened(const T &root,
                                                        MemoryPool *pool = nullptr) const {
            TURBO_MOVE_OR_RAISE(auto match, find_one_or_none(root));
            if (match.empty()) {
                return static_cast<get_type<T>>(nullptr);
            }
            return match.get_flattened(root, pool);
        }

    private:
        void Flatten(std::vector<FieldRef> children);

        std::variant<FieldPath, std::string, std::vector<FieldRef>> impl_;
    };

    TURBO_EXPORT void print_to(const FieldRef &ref, std::ostream *os);

    TURBO_EXPORT
    std::ostream &operator<<(std::ostream &os, const FieldRef &);

}  // namespace nebula
